package com.oak.broke.websocket;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.extra.spring.SpringUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.oak.broke.constant.WebSocketConstants;
import com.oak.broke.websocket.WebSocketSessions;
import jdk.nashorn.internal.runtime.GlobalConstants;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.AbstractWebSocketHandler;

import java.util.Map;
import java.util.UUID;

/**
 * 代理的第三方websocket接口handler
 */
@Slf4j
public class WebSocketRagClientProxyHandler extends AbstractWebSocketHandler {

    private final WebSocketSession webSocketServerSession;


    public WebSocketRagClientProxyHandler(WebSocketSession webSocketServerSession) {
        this.webSocketServerSession = webSocketServerSession;
    }

    /**
     * 第三方接口发送来的消息
     *
     * @param session
     * @param webSocketMessage
     * @throws Exception
     */
    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> webSocketMessage) throws Exception {

        if (StringUtils.equals(WebSocketConstants.PONG, (String) webSocketMessage.getPayload())) {
            return;
        }
        log.info("接收第三方socket接口返回，sessionId:{}，内容:{}", session.getId(), webSocketMessage.getPayload());

        Object payload = webSocketMessage.getPayload();
        JSONObject msgJson = JSONUtil.parseObj(payload);

        JSONObject header = msgJson.getJSONObject("header");
        JSONObject payloadInfo = msgJson.getJSONObject("payload");
        if (!"0".equals(header.getStr("code"))) {
            // 情况消息
            log.error("返回code错误");
        }


        String message = WebSocketSessions.RAG_WEBSOCKET_SESSIONS_MESSAGE_MAP.get(session.getId());
        StringBuffer sb = new StringBuffer(message == null ? "" : message);
        // 拼接ws返回的内容
        payloadInfo.getJSONObject("choices").getJSONArray("text").forEach(item -> {
            JSONObject text = JSONUtil.parseObj(item);
            //if (ModelConstants.ASSISTANT_ROLE.equals(text.getStr("role"))) {
            //    // 拼接对话内容
            //    sb.append(text.getStr("content"));
            //}
        });
        // 所有的回复内容
        WebSocketSessions.RAG_WEBSOCKET_SESSIONS_MESSAGE_MAP.put(session.getId(), sb.toString());

        String sid = WebSocketSessions.WEBSOCKET_SESSIONS_RAG_PROXY_MAP.get(session.getId());
        // 返回用户传入的问题参数
        Map wsReq = WebSocketSessions.RAG_WEBSOCKET_SESSIONS_SEND_MAP.get(sid);
        // 第一次返回，判断是文档还是模型
        if ("0".equals(header.getStr("status"))) {

        }
        // 最后一次回复
        if ("2".equals(header.getStr("status"))) {
            // 保存日志
            String content = WebSocketSessions.RAG_WEBSOCKET_SESSIONS_MESSAGE_MAP.get(session.getId());
            // todo 保存日志
            SpringUtil.getBean("xxxService");

            // 清除对话
            WebSocketSessions.RAG_WEBSOCKET_SESSIONS_MESSAGE_MAP.remove(session.getId());
        }
        // 收到第三方ws回复，转发给客户端
        // 匹配到的向量
        webSocketServerSession.sendMessage(new TextMessage(JSONUtil.toJsonStr(msgJson)));

    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        String sessionId = session.getId();
        log.info("清理session:status={},sessionId={}", status, sessionId);
        if (WebSocketSessions.WEBSOCKET_SESSIONS_RAG_PROXY_MAP.containsKey(sessionId)) {
            // 获取WS转发的会话id
            String proxySessionId = WebSocketSessions.WEBSOCKET_SESSIONS_RAG_PROXY_MAP.get(sessionId);
            if (StringUtils.isNotEmpty(proxySessionId)) {
                //关闭转发连接
                WebSocketSessions.WEBSOCKET_SESSIONS_RAG.get(proxySessionId).close();
            }
            WebSocketSessions.WEBSOCKET_SESSIONS_RAG_PROXY_MAP.remove(sessionId);
            log.info("清理session映射:sessionId={},proxySessionId={}", sessionId, proxySessionId);
        }
        //当前连接会话关闭时，清理session，同时关闭转发连接，清理session
        if (WebSocketSessions.WEBSOCKET_SESSIONS_RAG.containsKey(sessionId)) {
            session.close();
            log.info("当前连接会话关闭时，清理session，同时关闭转发连接，清理session");
            WebSocketSessions.WEBSOCKET_SESSIONS_RAG.remove(sessionId);
            WebSocketSessions.RAG_WEBSOCKET_SESSIONS_MESSAGE_MAP.remove(session.getId());
        }
    }
}
